#!/usr/bin/env python
"""
:mod:`nubio.io.fasta`
=====================

Provides read and write support for the fasta file format.

"""
from textwrap import wrap
from tempfile import NamedTemporaryFile
from nubox.char import CharVector, CharArray
from nubox.shim import izip

#def fasta_header_parser(ident):
# http://github.com/davisp/bioneb
#    bits = ident.split("|")
#    if len(bits) == 1:
#        return ident
#    ret = {}
#    while len(bits) > 0:
#        itype = bits[0]
#        parts = []
#        for b in bits[1:]:
#            if b in IDENT_TYPES: break
#            parts.append(b)
#        bits = bits[len(parts)+1:]
#        parts = filter(None, parts)
#        if len(parts) == 1:
#            parts = parts[0]
#        if isinstance(ret.get(itype, None), list):
#            ret[itype].append(parts)
#        elif itype in ret:
#            ret[itype] = [ret[itype], parts]
#        else:
#            ret[itype] = parts
#    return ret


def fasta_seq_parser(fh):
    """
    Parses a fasta file and generates ``tuples`` of ("sequence", "meta") for each
    sequence in the file. This parser does not preserve newlines in a fasta 
    file. The "sequence" is returned as ``bytes``, "meta" is a ``dict``. 
    
    If you want to create a ``CharVector`` (or sub-class thereof) instance you
    have to supply "sequence" and "meta" to its constructor 
    ``CharVector(*fasta_seq_parser(fh))``.
    
    The  "meta" dictionary corresponds to the "meta" argument of the default 
    ``NuBox`` constructor. Information specific to a fasta file are stored under
    the "fasta" key. Currently "comments", and "header" keys are created.
    
    Arguments:
    
        - fh(``file``) A file-like object open for reading with "fasta" content, 
          a sequence of lines is also valid.
    
    """
    header = None
    seqs = []
    comments = []
    meta = {}
    for line in fh:
        line = line.rstrip()
        if not line:
            continue
        if line[0] == ">":
            if header is not None:
                meta['fasta'] = {}
                meta['fasta']['header'] = header
                meta['fasta']['comments'] = comments
                yield "".join(seqs), meta
            header = line[1:].strip()
            seqs = []
            comments = []
            meta = {}
        elif line[0] == ';':
            # unusual comment
            comments.append(line[1:].strip())
        else:
            seqs.append(line)
    meta['fasta'] = {}
    meta['fasta']['header'] = header
    meta['fasta']['comments'] = comments
    yield "".join(seqs), meta

def fasta_seq_writer(fh, vector, width=80, **kwargs):
    """
    Writes a ``CharVector`` (or sub-class thereof) instance or a ("sequence", 
    "meta") ``tuple`` to an open for writing ``file`` in the fasta file 
    format. Supports other file-like objects with a ``write`` method. The width
    of the sequence (not comments) can be adjusted with the "width" argument. 
    The "meta" dictionary or ``vector.meta`` property must include 
    ``meta["fasta"]["header"]``.
    
    The  "meta" dictionary should corresponds to the "meta" argument of the 
    default  ``NuBox`` constructor. Information specific to a fasta file are 
    stored under the "fasta" key. Currently "comments", and "header" keys are 
    created.
    
    Arguments:
    
        - fh(``file`` or file-like) an open for writing file handle or 
          file-like object with a ``write`` method.
        - vector(``CharVector`` or sequence) if sequence the contents should be
          "letters" and "meta", where "letters" is a sequence of  ``bytes`` 
          and comments is a ``dict``.
    
    Additional keyworded arguments are passed to the ``textwrap.wrap`` function.
    
    """
    # get data
    if isinstance(vector, CharVector):
        try:
            seq = vector.data.tostring()
        except AttributeError:
            seq = str(vector.data)
        meta = vector.meta
    else:
        seq, meta = vector

    header = meta.get('fasta', {}).get('header', '')
    comments = meta.get('fasta', {}).get('comments', ())
    # write data
    lines = [">%s" % header]
    lines.extend([';%s' % line for line in comments])
    lines.extend(wrap(seq, width, **kwargs))
    fh.write("\n".join(lines))
    fh.write("\n")

def fasta_aln_parser(fh):
    """
    Parses a fasta file with multiple sequences and generates two lists one 
    of the "sequences" and one of "meta" associated with the "sequences". The
    "sequences" are ``bytes``. The "meta" dictionary corresponds to the "meta"
    argument of the default ``NuBox`` constructor. Information specific to a
    fasta file are stored under the "fasta" key. Currently "comments", and 
    "header" keys are created.
    
    Arguments:
    
        - fh(``file``) an open for reading file-like with fasta
          content, a sequence of lines or a ``StringIO`` object.
    
    """
    parser = fasta_seq_parser(fh)
    seqs = []
    seq_metas = []
    for seq, meta in parser:
        seqs.append(seq)
        seq_metas.append(meta)
    # no
    aln_meta = {'fasta':{}}
    return seqs, seq_metas, aln_meta

def fasta_aln_writer(fh, vector, width=80, **kwargs):
    """
    Writes a ``SymbolArray`` instance or a (sequences, metadata) tuple to an 
    open for writing file handle in the fasta file format. Supports other 
    objects with a ``write`` method. The "width" of the sequence can be 
    specified. Sequences should be valid ``fasta_seq_writer`` vectors.
    Does not support alignment metadata. Only supports sequences comments.
    Returns ``None``.
    
    Arguments:
    
        - fh(``file``) An open for writing file handle or object with a 
          ``write`` method.
        - vector(``CharVector`` instance or sequence) The contents of the 
          "sequence" should be (sequences, metadata), where sequences is a valid
          input for ``fasta_seq_writer``. And metadata is the (currently) 
          ignored per-alignement metadata.
          
    Additional keyworded arguments are passed to the ``textwrap.wrap`` function.
    
    """
    if isinstance(vector, CharArray):
        for seq in vector:
            # this could be done without child iteration (faster)
            fasta_seq_writer(fh, seq, width, **kwargs)
    else:
        for seq, seq_meta in izip(*vector[0:2]):
            if 'meta' in kwargs:
                seq_meta.update(kwargs['meta'])
            fasta_seq_writer(fh, (seq, seq_meta), width, **kwargs)

def fasta_tmp(vector, **kwargs):
    """
    Writes sequence(s) to a temporary file. Returns a file handle open for 
    reading and writing. If the handle is closed or garbage collected the 
    temporary file will be deleted from the file system. Contents of the file 
    will be in the fasta file format.
    
    Arguments:
    
        - vector(``CharVector`` instance or sequence). This should be a valid 
          input for ``fasta_writer`` or a ``CharArray`` if multiple sequences
          are to be written to one file.
          
    Additional keyworded arguments are passed to ``fasta_seq_writer``.
    
    """
    tmp_file = NamedTemporaryFile()
    is_vector = isinstance(vector, CharVector)
    is_array = isinstance(vector, CharArray)
    if is_array:
        fasta_aln_writer(tmp_file, vector, **kwargs)
    elif is_vector:
        fasta_seq_writer(tmp_file, vector, **kwargs)
    else:
        # got seq meta or seqs metas
        if isinstance(vector[1], dict):
            # seq str, meta dict
            fasta_seq_writer(tmp_file, vector, **kwargs)
        else:
            # seqs (list) metas (list)
            fasta_aln_writer(tmp_file, vector, **kwargs)
    tmp_file.flush()
    return tmp_file
